#include "zen_image.h"
#include "zen_random.h"
#include "zen_file.h"
#include "zen_image_coder.h"
#include "zen_image_filter.h"
#include "zen_image_stb.h"
#include "zen_number.h"
#include "zen_endian.h"
#include <iostream>
#include <vector>

std::string bin;
void usage()
{
	std::cout << bin << " -e from_file to_file code" << std::endl;
	std::cout << bin << " -p from_file to_image_file" << std::endl;
	std::cout << bin << " -r from_image_file to_file" << std::endl;
	exit(1);
}
void error(std::string const& e)
{
	std::cout << e << std::endl;
	exit(2);
}

//////////////////////// ////////////////////////
int get_min_edge(int sz)
{
	return (int)sqrt((sz - 1) / 3) + 1;
}
void RevertImage(std::string from, std::string to)
{
	auto source = Zen::File::LoadFileToBuffer(from);
	Zen::ImageSTBDecoder dec;
	auto image = dec.decode(source);
	auto data = image->Data();
	if (image->Size() < 4)
	{
		error("wrong image size");
	}
	Zen::Byte4 b;
	auto p = (data + image->Size() - 4);
	b.Set(p[0], p[1], p[2], p[3]);
	auto size = b.GetBEValue();
	if (size > image->Size() - 4)
	{
		error("wrong image size");
	}
	Zen::File::WriteBufferToFile(to, data, size);
}

void ToImage(std::string from, std::string to)
{
	auto source = Zen::File::LoadFileToBuffer(from);
	if (source.empty())
	{
		error("empty file:" + from);
	}

	int sz = (int)source.size();

	auto n = get_min_edge(sz + 4);

	auto new_size = n * n * 3;
	source.resize(new_size, 0);

	Zen::Byte4 b;
	b.SetWithBE(sz);
	auto tail = (source.data() + source.size() - 4);
	for (int i = 0; i < 4; ++i) tail[i] = b.bytes[i];

	auto image = Zen::Image::GenerateWithData(Zen::EPixel::RGB, n, n, source.data());

	Zen::ImageBMPEncoder png;
	auto data = png.encode(*image);

	Zen::File::WriteBufferToFile(to, data);
}

//////////////////////// ////////////////////////
void encode(std::vector<uint8_t>& buffer, int code)
{
	Zen::Random random(code);
	for (auto& v : buffer)
	{
		v ^= (uint8_t)(random() % 256);
	}
}

void run_encode(std::string from, std::string to, int code)
{
	std::cout << "code:" << code << std::endl;
	std::cout << "from file:" << from << std::endl;
	std::cout << "to file:" << to << std::endl;

	auto data = Zen::File::LoadFileToBuffer(from);
	if (data.empty())
	{
		error("no data");
	}

	encode(data, code);
	if (!Zen::File::WriteBufferToFile(to, data))
	{
		error("save failed");
	}
}
//////////////////////// ////////////////////////
int main(int argn, char const** args)
{
	if (argn > 0)
	{
		bin = Zen::PickFileName(args[0]);
	}
	if (argn < 2)
	{
		usage();
	}

	std::string cmd = args[1];
	if (cmd == "-e")
	{
		int code = 0;
		if (argn < 4) usage();
		if (argn > 4) code = atoi(args[4]);
		run_encode(args[2], args[3], code);
	}
	else if (cmd == "-p")
	{
		if (argn != 4) usage();
		ToImage(args[2], args[3]);
	}
	else if (cmd == "-r")
	{
		if (argn != 4) usage();
		RevertImage(args[2], args[3]);
	}
	return 0;
}
